import { Express, NextFunction, Request, Response } from 'express'

// XSS防护
const xssProtection = (req: Request, _res: Response, next: NextFunction) => {
  if (req.body) {
    for (let key in req.body) {
      if (typeof req.body[key] === 'string') {
        req.body[key] = req.body[key]
          .replace(/</g, '&lt;')
          .replace(/>/g, '&gt;')
          .replace(/"/g, '&quot;')
          .replace(/'/g, '&#x27;')
          .replace(/\//g, '&#x2F;')
      }
    }
  }
  next()
}

// NoSQL注入防护
const noSqlInjectionProtection = (req: Request, _res: Response, next: NextFunction) => {
  if (req.body) {
    const sanitize = (obj: any): any => {
      if (Array.isArray(obj)) {
        return obj.map(sanitize)
      }
      if (obj && typeof obj === 'object') {
        const result: any = {}
        for (const key in obj) {
          if (key.startsWith('$')) {
            // 跳过MongoDB操作符
            continue
          }
          result[key] = sanitize(obj[key])
        }
        return result
      }
      return obj
    }
    req.body = sanitize(req.body)
  }
  next()
}

// 安全头部设置
const securityHeaders = (_req: Request, res: Response, next: NextFunction) => {
  // 基本安全头
  res.setHeader('X-Content-Type-Options', 'nosniff')
  res.setHeader('X-Frame-Options', 'SAMEORIGIN')
  res.setHeader('X-XSS-Protection', '1; mode=block')
  res.setHeader('Referrer-Policy', 'same-origin')
  res.setHeader('Content-Security-Policy', "default-src 'self'")

  // 删除不安全的头
  res.removeHeader('X-Powered-By')

  next()
}

export const setupSecurity = (app: Express) => {
  // CORS设置
  app.use((_req, res, next) => {
    res.setHeader('Access-Control-Allow-Origin', process.env.CLIENT_URL || '*')
    res.setHeader('Access-Control-Allow-Methods', 'GET, POST, PUT, DELETE, PATCH, OPTIONS')
    res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization')
    next()
  })

  // 应用安全中间件
  app.use(xssProtection)
  app.use(noSqlInjectionProtection)
  app.use(securityHeaders)
}
